#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
主窗口界面
"""

import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')

from PyQt5.QtWidgets import QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QTextEdit, QLabel, QFileDialog, QComboBox, QMessageBox, QProgressBar, QApplication, QListWidget, QSplitter, QScrollArea, QGridLayout, QListWidgetItem
from PyQt5.QtCore import QThread, pyqtSignal, QSize
from PyQt5.QtGui import QPixmap, QImage, QIcon

# 导入日志模块
from utils.logger import get_logger

# 安全导入模块，提供错误处理
PDF_PROCESSOR_AVAILABLE = False
PDFProcessor = None

OCR_ENGINE_AVAILABLE = False
OCREngine = None

EXPORTERS_AVAILABLE = False
PDFExporter = None
DOCXExporter = None
MDExporter = None

# 创建日志记录器
logger = get_logger('MainWindow')

try:
    from ocr.pdf_processor import PDFProcessor as PDFProcessorClass
    PDFProcessor = PDFProcessorClass
    PDF_PROCESSOR_AVAILABLE = True
    logger.info("PDF处理器导入成功")
except Exception as e:
    error_msg = f"PDF处理器导入失败: {str(e)}"
    logger.error(error_msg)
    print(error_msg)

try:
    from ocr.ocr_engine import OCREngine as OCREngineClass
    OCREngine = OCREngineClass
    OCR_ENGINE_AVAILABLE = True
    logger.info("OCR引擎导入成功")
except Exception as e:
    error_msg = f"OCR引擎导入失败: {str(e)}"
    logger.error(error_msg)
    print(error_msg)

try:
    from export.pdf_exporter import PDFExporter as PDFExporterClass
    from export.docx_exporter import DOCXExporter as DOCXExporterClass
    from export.md_exporter import MDExporter as MDExporterClass
    PDFExporter = PDFExporterClass
    DOCXExporter = DOCXExporterClass
    MDExporter = MDExporterClass
    EXPORTERS_AVAILABLE = True
    logger.info("导出模块导入成功")
except Exception as e:
    error_msg = f"导出模块导入失败: {str(e)}"
    logger.error(error_msg)
    print(error_msg)

class WorkerThread(QThread):
    """工作线程，用于执行耗时的OCR操作"""
    finished = pyqtSignal(dict)  # 修改为发送字典，包含图片路径和识别结果
    progress = pyqtSignal(int)
    error = pyqtSignal(str)

    def __init__(self, image_paths, parent=None):
        super().__init__(parent)
        self.image_paths = image_paths
        self.logger = get_logger('WorkerThread')
        # 安全初始化OCR引擎
        try:
            if OCR_ENGINE_AVAILABLE and OCREngine is not None:
                # 初始化OCR引擎
                self.ocr_engine = OCREngine()
                self.logger.info("OCR引擎初始化成功")
            else:
                self.ocr_engine = None
                self.logger.error("OCR引擎不可用")
        except Exception as e:
            self.ocr_engine = None
            error_msg = f"OCR引擎初始化失败: {str(e)}"
            self.logger.error(error_msg)
            print(error_msg)

    def run(self):
        # 检查OCR引擎是否可用
        if self.ocr_engine is None:
            error_msg = "OCR引擎不可用，请检查安装"
            self.error.emit(error_msg)
            self.logger.error(error_msg)
            return
            
        try:
            results = {}
            detailed_results = {}  # 保存详细结果用于原样排版
            total = len(self.image_paths)
            self.logger.info(f"开始OCR识别，共{total}张图片")
            
            for i, image_path in enumerate(self.image_paths):
                self.logger.info(f"正在处理图片: {image_path}")
                # 执行OCR识别
                result = self.ocr_engine.get_detailed_results(image_path)
                results[image_path] = result
                
                # 发送进度信号
                progress = int((i + 1) / total * 100)
                self.progress.emit(progress)
            
            # 发送所有结果
            self.logger.info("OCR识别完成")
            self.finished.emit({"results": results, "detailed_results": detailed_results})
        except Exception as e:
            error_msg = f"OCR识别过程中发生错误: {str(e)}"
            self.logger.error(error_msg)
            self.error.emit(error_msg)

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.logger = get_logger('MainWindow')
        self.logger.info("主窗口初始化开始")
        self.setWindowTitle("PaddleOCR PDF文字识别工具")
        self.setGeometry(100, 100, 1200, 800)
        
        # 当前处理的PDF路径
        self.current_pdf_path = None
        self.image_paths = []  # 存储PDF转换后的图片路径
        
        # 初始化UI
        self.init_ui()
        
        # 初始化导出器
        if EXPORTERS_AVAILABLE and PDFExporter is not None and DOCXExporter is not None and MDExporter is not None:
            self.pdf_exporter = PDFExporter()
            self.docx_exporter = DOCXExporter()
            self.md_exporter = MDExporter()
            self.logger.info("文档导出器初始化成功")
        else:
            self.pdf_exporter = None
            self.docx_exporter = None
            self.md_exporter = None
            self.logger.warning("文档导出器初始化失败")
        
        # 初始化OCR引擎
        self.worker_thread = None
        self.logger.info("主窗口初始化完成")
    
    def init_ui(self):
        self.logger.info("UI初始化开始")
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QVBoxLayout()
        central_widget.setLayout(main_layout)
        
        # 创建顶部按钮区域
        top_layout = QHBoxLayout()
        
        self.open_button = QPushButton("打开PDF")
        self.open_button.clicked.connect(self.open_pdf)
        
        self.ocr_button = QPushButton("开始OCR识别")
        self.ocr_button.clicked.connect(self.start_ocr)
        self.ocr_button.setEnabled(False)
        
        self.save_button = QPushButton("保存")
        self.save_button.clicked.connect(self.save_document)
        self.save_button.setEnabled(False)
        
        # 添加原样排版保存按钮
        self.save_layout_button = QPushButton("保存(原样排版)")
        self.save_layout_button.clicked.connect(self.save_document_with_layout)
        self.save_layout_button.setEnabled(False)
        
        self.format_combo = QComboBox()
        self.format_combo.addItem("PDF")
        self.format_combo.addItem("DOCX")
        self.format_combo.addItem("Markdown")
        
        top_layout.addWidget(self.open_button)
        top_layout.addWidget(self.ocr_button)
        top_layout.addWidget(QLabel("保存格式:"))
        top_layout.addWidget(self.format_combo)
        top_layout.addWidget(self.save_button)
        top_layout.addWidget(self.save_layout_button)
        top_layout.addStretch()
        
        # 创建进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        
        # 创建分割器以实现两栏布局
        splitter = QSplitter()
        
        # 左侧：图片列表区域
        left_widget = QWidget()
        left_layout = QVBoxLayout()
        left_widget.setLayout(left_layout)
        
        left_label = QLabel("PDF转换后的图片")
        left_layout.addWidget(left_label)
        
        self.image_list = QListWidget()
        self.image_list.setIconSize(QSize(200, 200))  # 设置图标大小
        self.image_list.clicked.connect(self.on_image_selected)
        left_layout.addWidget(self.image_list)
        
        # 右侧：识别结果显示区域
        right_widget = QWidget()
        right_layout = QVBoxLayout()
        right_widget.setLayout(right_layout)
        
        right_label = QLabel("OCR识别结果")
        right_layout.addWidget(right_label)
        
        # 创建文本编辑区域
        self.text_edit = QTextEdit()
        self.text_edit.setPlaceholderText("OCR识别结果显示在这里，您可以编辑识别结果...")
        right_layout.addWidget(self.text_edit)
        
        # 将左右两侧添加到分割器
        splitter.addWidget(left_widget)
        splitter.addWidget(right_widget)
        splitter.setSizes([400, 800])  # 设置初始大小比例
        
        # 添加到主布局
        main_layout.addLayout(top_layout)
        main_layout.addWidget(self.progress_bar)
        main_layout.addWidget(splitter)
        self.logger.info("UI初始化完成")
    
    def open_pdf(self):
        self.logger.info("打开PDF文件")
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择PDF文件", "", "PDF Files (*.pdf)"
        )
        
        if file_path:
            self.logger.info(f"选择的PDF文件: {file_path}")
            self.current_pdf_path = file_path
            self.ocr_button.setEnabled(True)
            self.image_list.clear()  # 清空图片列表
            self.text_edit.clear()   # 清空文本编辑区域
            
            # 立即转换PDF为图片并显示在左栏
            self.convert_pdf_to_images()
            
            status_bar = self.statusBar()
            if status_bar is not None:
                status_bar.showMessage(f"已选择文件: {os.path.basename(file_path)}")
    
    def convert_pdf_to_images(self):
        """将PDF转换为图片并显示在左栏"""
        self.logger.info("开始转换PDF为图片")
        if not self.current_pdf_path:
            self.logger.warning("未选择PDF文件")
            return
            
        if not PDF_PROCESSOR_AVAILABLE or PDFProcessor is None:
            self.logger.error("PDF处理器不可用")
            return
            
        try:
            # 创建PDF处理器
            pdf_processor = PDFProcessor()
            
            # 提取PDF页面为图像
            self.image_paths = pdf_processor.extract_pages_as_images(self.current_pdf_path)
            self.logger.info(f"PDF转换完成，共生成{len(self.image_paths)}张图片")
            
            # 在左栏显示图片列表
            self.display_images_in_list()
            
        except Exception as e:
            error_msg = f"转换PDF为图片时发生错误: {str(e)}"
            self.logger.error(error_msg)
            QMessageBox.critical(self, "错误", error_msg)
    
    def display_images_in_list(self):
        """在左栏列表中显示图片"""
        self.logger.info("开始在列表中显示图片")
        self.image_list.clear()
        for i, image_path in enumerate(self.image_paths):
            # 创建列表项
            item = QListWidgetItem(f"页面 {i+1}")
            
            # 加载并设置图片缩略图
            pixmap = QPixmap(image_path)
            if not pixmap.isNull():
                # 缩放图片以适应列表项
                try:
                    scaled_pixmap = pixmap.scaled(200, 200)
                    # 使用QIcon包装QPixmap
                    icon = QIcon(scaled_pixmap)
                    item.setIcon(icon)
                except Exception as e:
                    # 如果缩放失败，使用原始图片
                    self.logger.warning(f"图片缩放失败: {str(e)}")
                    icon = QIcon(pixmap)
                    item.setIcon(icon)
            
            # 保存图片路径作为数据 (使用数值替代Qt.UserRole)
            try:
                item.setData(32, image_path)  # UserRole = 32
            except Exception as e:
                self.logger.warning(f"设置图片路径数据失败: {str(e)}")
            self.image_list.addItem(item)
        self.logger.info("图片显示完成")
    
    def on_image_selected(self, index):
        """当用户选择图片时的处理函数"""
        # 这里可以实现点击图片时显示对应的OCR结果
        # 暂时留空，后续可以扩展
        self.logger.info(f"选择了图片索引: {index.row()}")
        pass
    
    def start_ocr(self):
        self.logger.info("开始OCR识别")
        if not self.image_paths:
            warning_msg = "请先打开PDF文件"
            self.logger.warning(warning_msg)
            QMessageBox.warning(self, "提示", warning_msg)
            return
            
        # 检查依赖是否可用
        if not OCR_ENGINE_AVAILABLE or OCREngine is None:
            error_msg = "OCR引擎不可用，请检查安装"
            self.logger.error(error_msg)
            QMessageBox.critical(self, "错误", error_msg)
            return
            
        try:
            # 创建并启动工作线程
            self.worker_thread = WorkerThread(self.image_paths)
            self.worker_thread.finished.connect(self.on_ocr_finished)
            self.worker_thread.progress.connect(self.on_progress_updated)
            self.worker_thread.error.connect(self.on_ocr_error)
            
            # 禁用按钮，显示进度条
            self.ocr_button.setEnabled(False)
            self.open_button.setEnabled(False)
            self.progress_bar.setVisible(True)
            self.progress_bar.setValue(0)
            
            # 开始OCR识别
            self.worker_thread.start()
            self.logger.info("OCR识别线程已启动")
            
        except Exception as e:
            error_msg = f"启动OCR识别时发生错误: {str(e)}"
            self.logger.error(error_msg)
            QMessageBox.critical(self, "错误", error_msg)
    
    def on_ocr_finished(self, results):
        self.logger.info("OCR识别完成")
        # 保存OCR结果用于原样排版
        if isinstance(results, dict) and "results" in results:
            self.ocr_results = results["results"]  # 保存按页面分组的结果
            
            # 更新UI - 显示识别结果
            full_text = ""
            for image_path, page_results in results["results"].items():
                # 添加页面标题
                page_num = self.image_paths.index(image_path) + 1
                full_text += f"\n========== 页面 {page_num} ==========\n"
                
                # 添加该页面的识别结果
                for result in page_results:
                    if isinstance(result, dict) and "text" in result:
                        full_text += result["text"] + "\n"
                    else:
                        full_text += str(result) + "\n"
        else:
            # 兼容旧格式
            self.ocr_results = results  # 保存按页面分组的结果
            
            # 更新UI - 显示识别结果
            full_text = ""
            for image_path, page_results in results.items():
                # 添加页面标题
                page_num = self.image_paths.index(image_path) + 1
                full_text += f"\n========== 页面 {page_num} ==========\n"
                
                # 添加该页面的识别结果
                for result in page_results:
                    if isinstance(result, dict) and "text" in result:
                        full_text += result["text"] + "\n"
                    else:
                        full_text += str(result) + "\n"
        
        # 更新文本编辑区域
        self.text_edit.setPlainText(full_text)
        self.save_button.setEnabled(True)
        self.save_layout_button.setEnabled(True)  # 启用原样排版保存按钮
        self.progress_bar.setVisible(False)
        self.ocr_button.setEnabled(True)
        self.open_button.setEnabled(True)
        status_bar = self.statusBar()
        if status_bar is not None:
            status_bar.showMessage("OCR识别完成")
    
    def on_progress_updated(self, value):
        self.progress_bar.setValue(value)
    
    def on_ocr_error(self, error_msg):
        self.logger.error(f"OCR错误: {error_msg}")
        QMessageBox.critical(self, "OCR错误", f"OCR识别过程中发生错误: {error_msg}")
        self.progress_bar.setVisible(False)
        self.ocr_button.setEnabled(True)
        self.open_button.setEnabled(True)
    
    def save_document(self):
        self.logger.info("开始保存文档")
        if not self.current_pdf_path:
            self.logger.warning("未选择PDF文件")
            return
            
        # 检查导出器是否可用
        if not EXPORTERS_AVAILABLE or self.pdf_exporter is None or self.docx_exporter is None or self.md_exporter is None:
            error_msg = "文档导出功能不可用，请检查安装"
            self.logger.error(error_msg)
            QMessageBox.critical(self, "错误", error_msg)
            return
            
        # 获取文本内容
        content = self.text_edit.toPlainText()
        
        # 根据选择的格式确定文件扩展名
        format_selected = self.format_combo.currentText()
        extension_map = {
            "PDF": ".pdf",
            "DOCX": ".docx",
            "Markdown": ".md"
        }
        extension = extension_map.get(format_selected, ".txt")
        
        # 构造默认保存路径
        base_name = os.path.splitext(os.path.basename(self.current_pdf_path))[0]
        default_filename = f"{base_name}_ocr{extension}"
        
        # 打开保存对话框
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存文件", default_filename, f"{format_selected} Files (*{extension})"
        )
        
        if file_path:
            try:
                self.logger.info(f"保存文件到: {file_path}")
                # 根据选择的格式保存文件
                if format_selected == "PDF":
                    self.save_as_pdf(file_path, content)
                elif format_selected == "DOCX":
                    self.save_as_docx(file_path, content)
                elif format_selected == "Markdown":
                    self.save_as_markdown(file_path, content)
                
                success_msg = f"文件已保存到: {file_path}"
                self.logger.info(success_msg)
                QMessageBox.information(self, "成功", success_msg)
                status_bar = self.statusBar()
                if status_bar is not None:
                    status_bar.showMessage(f"文件已保存: {file_path}")
            except Exception as e:
                error_msg = f"保存文件时发生错误: {str(e)}"
                self.logger.error(error_msg)
                QMessageBox.critical(self, "保存失败", error_msg)
    
    def save_document_with_layout(self):
        """保存文档（保持原样排版）"""
        self.logger.info("开始保存文档（原样排版）")
        if not self.current_pdf_path:
            self.logger.warning("未选择PDF文件")
            return
            
        # 检查导出器是否可用
        if not EXPORTERS_AVAILABLE or self.pdf_exporter is None or self.docx_exporter is None:
            error_msg = "文档导出功能不可用，请检查安装"
            self.logger.error(error_msg)
            QMessageBox.critical(self, "错误", error_msg)
            return
            
        # 根据选择的格式确定文件扩展名
        format_selected = self.format_combo.currentText()
        extension_map = {
            "PDF": ".pdf",
            "DOCX": ".docx",
            "Markdown": ".md"
        }
        extension = extension_map.get(format_selected, ".txt")
        
        # 构造默认保存路径
        base_name = os.path.splitext(os.path.basename(self.current_pdf_path))[0]
        default_filename = f"{base_name}_ocr_layout{extension}"
        
        # 打开保存对话框
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存文件(原样排版)", default_filename, f"{format_selected} Files (*{extension})"
        )
        
        if file_path:
            try:
                self.logger.info(f"保存文件到: {file_path}")
                # 根据选择的格式保存文件（保持原样排版）
                if format_selected == "PDF":
                    self.save_as_layout_preserved_pdf(file_path)
                elif format_selected == "DOCX":
                    self.save_as_layout_preserved_docx(file_path)
                elif format_selected == "Markdown":
                    # Markdown格式暂时使用普通导出
                    content = self.text_edit.toPlainText()
                    self.save_as_markdown(file_path, content)
                
                success_msg = f"文件已保存到: {file_path}"
                self.logger.info(success_msg)
                QMessageBox.information(self, "成功", success_msg)
                status_bar = self.statusBar()
                if status_bar is not None:
                    status_bar.showMessage(f"文件已保存: {file_path}")
            except Exception as e:
                error_msg = f"保存文件时发生错误: {str(e)}"
                self.logger.error(error_msg)
                QMessageBox.critical(self, "保存失败", error_msg)
    
    def save_as_pdf(self, file_path, content):
        """保存为PDF格式（支持中英文）"""
        self.logger.info("保存为PDF格式")
        if self.pdf_exporter is not None:
            self.pdf_exporter.export_text_to_pdf(file_path, content)
    
    def save_as_docx(self, file_path, content):
        """保存为DOCX格式"""
        self.logger.info("保存为DOCX格式")
        if self.docx_exporter is not None:
            self.docx_exporter.export_text_to_docx(file_path, content)
    
    def save_as_markdown(self, file_path, content):
        """保存为Markdown格式"""
        self.logger.info("保存为Markdown格式")
        if self.md_exporter is not None:
            self.md_exporter.export_text_to_md(file_path, content)
    
    def save_as_layout_preserved_pdf(self, file_path):
        """保存为保持原样排版的PDF格式"""
        self.logger.info("保存为保持原样排版的PDF格式")
        if self.pdf_exporter is not None and hasattr(self, 'ocr_results'):
            self.pdf_exporter.export_layout_preserved_pdf(file_path, self.ocr_results)
    
    def save_as_layout_preserved_docx(self, file_path):
        """保存为保持原样排版的DOCX格式"""
        self.logger.info("保存为保持原样排版的DOCX格式")
        if self.docx_exporter is not None and hasattr(self, 'ocr_results'):
            # 对于DOCX，我们按页面分组处理
            try:
                from docx import Document
                doc = Document()
                doc.add_heading('OCR识别结果（原样排版）', 0)
                
                # 按页面处理
                for image_path, page_results in self.ocr_results.items():
                    # 添加页面标题
                    page_num = self.image_paths.index(image_path) + 1
                    doc.add_heading(f'页面 {page_num}', 1)
                    
                    # 使用布局保持器处理单页
                    if hasattr(self.docx_exporter, 'layout_preserver'):
                        page_content = self.docx_exporter.layout_preserver.preserve_layout_from_page_results(page_results, page_num, "text")
                    else:
                        # 简单处理
                        page_content = "\n".join([result["text"] for result in page_results if isinstance(result, dict) and "text" in result])
                    
                    # 添加页面内容
                    paragraph = doc.add_paragraph()
                    paragraph.add_run(page_content)
                
                doc.save(file_path)
            except Exception as e:
                self.logger.error(f"保存DOCX时发生错误: {str(e)}")
                raise
